import numpy

####################################################################
# optimize()
# OVERVIEW: Runs a simple genetic algorithm to optimize an MCSP
# PARAMETERS:
#       siteSppMatrix, the 2D, binary numpy array to shuffle (rows are sites, columns are species)
#           Alternatively, siteSppMatrix can be a file path for an ASCII file containing the matrix
#       popSize, the number of solutions to create/test each generation
#       nElite, the number of best solutions to keep without mutation each generation
#       nGenerations, the number of times to loop through the GA (unless convergence)
# RETURNS: the set of sites that should be protected
####################################################################
def optimize(siteSppMatrix,siteBudget,popSize=100,nElite=10,nGenerations=1000):
    # Load data from the file if applicable
    if type(siteSppMatrix) == type('s'): siteSppMatrix = numpy.loadtxt(siteSppMatrix)
   
    # Convert to boolean for faster computation (SHOULD THIS BE BOOLEAN?)
    siteSppMatrix = siteSppMatrix.astype(bool)
    
    # Other initial stuff (parameters, etc)
    crossRate = 0.6 # probability that two parents chosen will crossover vs create identical offspring
    mutRate = 0.1 # probability that an offspring mutates
    nSites,nSpecies = siteSppMatrix.shape
    
    # Define parameters for early termination of the optimization
    stopCheckInterval = 0.1 # See stopSlope explanation
    stopCV = 0.001 # Coefficient of variation (std_dev/mean) of the max fitnesses over 
        # (stopCheckInterval*nGenerations) of the iterations to stop the execution early. For instance, 
        # if nGenerations = 100 and stopCheckInterval = 0.1, then this would take the maximum fitnesses 
        # of the previous 10 generations, and divide their standard deviation by their average. If the 
        # resulting coefficient of variation is less than stopCV the algorithm would stop early. Setting
        # this value stops the algorithm from running when its not improving.

    ############## I. CREATE THE INITIAL POPULATION #################
    # Create the initial population (solution set), which is composed
    #   of sets of sites that are to be protected
    currentPopulation = numpy.array([numpy.random.permutation(nSites)[:siteBudget] for i in xrange(popSize)])
    
    # Get the fitness (number of species represented w/ complementarity included) of the population
    currentFitnesses = numpy.array([siteSppMatrix[currentPopulation[i]].any(axis=0).sum() for i in xrange(popSize)])
    
    # Do the dumb check to see if only 1 site is in the budget
    if siteBudget == 1:
        randArrangement = numpy.random.permutation(popSize)
        currentPopulation = currentPopulation[randArrangement]
        currentFitnesses = currentFitnesses[randArrangement]
        bestIndex = numpy.argmax(currentFitnesses)
        bestSites = currentPopulation[bestIndex]
        speciesCovered = numpy.arange(nSpecies)[siteSppMatrix[currentPopulation[bestIndex]].any(axis=0)]
        print 'Execution ended early because only 1 site was in the budget, so no optimization is necessary.'
        
        return (bestSites,speciesCovered)
        
    # Define iteration variables
    allSites = numpy.arange(nSites)
    maxFitnesses = numpy.zeros(nGenerations,dtype=int)
    stopCheckIntervalNumber = round(stopCheckInterval*nGenerations)
    
    # Start the iteration
    for g in xrange(nGenerations):
    
        # Determine if the algorithm should stop early
        maxFitnesses[g] = currentFitnesses.max()
        if maxFitnesses[g] == nSpecies: # Then all species are covered
            print 'Execution ended early on generation %i because all species are covered.' % g
            break
                
        elif g >= stopCheckIntervalNumber:
            fitnessVec = maxFitnesses[g-stopCheckIntervalNumber:g]
            fitnessCV = fitnessVec.std()/fitnessVec.mean()
            if fitnessCV < stopCV:
                print 'Execution ended early on generation %i because the ' % g + \
                    'coefficient of varation of maximum fitnesses over ' + \
                    '%0.0f generations is less than %0.4f.' % (stopCheckIntervalNumber, stopCV)
                break
        
        # Sort the current population, descending by fitness
        fitnessSort = numpy.argsort(currentFitnesses)[::-1]
        currentFitnesses = currentFitnesses[fitnessSort]
        currentPopulation = currentPopulation[fitnessSort]
    
        nextPopulation = numpy.zeros((popSize,siteBudget),dtype=int)
        nextFitnesses = numpy.zeros(popSize,dtype=int)
                
        ############## II. CHOOSE THE BEST SOLUTIONS TO KEEP WITHOUT MUTATION #################
        nextPopulation[:nElite] = currentPopulation[:nElite]
        nextFitnesses[:nElite] = currentFitnesses[:nElite]
        
        for i in xrange(nElite,popSize):
            ############## III. CHOOSE THE PARENTS (SOLUTIONS TO CROSS) #################
            parentSiteInds = numpy.argsort(numpy.random.rand(popSize)*currentFitnesses)[-2:]
            
            ############ IV. CROSSOVER (MATING) ##############
            # Determine if the cross happens
            if numpy.random.rand() <= crossRate: # Do the cross
                parentSites = currentPopulation[parentSiteInds]
                siteUnion = numpy.union1d(parentSites[0],parentSites[1])
                offspringSites = numpy.random.permutation(siteUnion)[:siteBudget]
                
            else: #offspring is identical to one parent (before mutation)
                numpy.random.shuffle(parentSiteInds)
                offspringSites = currentPopulation[parentSiteInds[0]]
                    
            if numpy.unique(offspringSites).size < offspringSites.size: print 'repeat sites!'
            
            ############## V. MUTATION ##############
            # Mutate the offspring's row and column orderings with probability mutRate
            if numpy.random.rand() <= mutRate:
                # Decide the number of sites to change in the offspring solution
                nSwaps = numpy.random.random_integers(0,siteBudget)
                
                # Choose which sites will be changed and exchange them with unused sites in the total set
                offspringSites = numpy.random.permutation(offspringSites)
                
                unusedSites = numpy.setdiff1d(allSites,offspringSites)
                
                if unusedSites.size < nSwaps:
                    unusedSites = numpy.concatenate((unusedSites,offspringSites[:nSwaps])) # Guarantees there will be enough sites to swap out
                
                swapSites = numpy.random.permutation(unusedSites)
                offspringSites[:nSwaps] = swapSites[:nSwaps]
                
            if numpy.unique(offspringSites).size < offspringSites.size: print 'repeat sites!'
            
            # Add the offspring to the next generation's population
            nextPopulation[i] = offspringSites
            nextFitnesses[i] = siteSppMatrix[offspringSites].any(axis=0).sum() # Fitness of the offspring
            
        # Update the population to the newest generation
        currentPopulation = nextPopulation
        currentFitnesses = nextFitnesses
    
    # Return the set of sites with the highest fitness
    bestIndex = numpy.argmax(currentFitnesses)
    bestSites = numpy.sort(currentPopulation[bestIndex])
    speciesCovered = numpy.arange(nSpecies)[siteSppMatrix[currentPopulation[bestIndex]].any(axis=0)]
        
    return (bestSites,speciesCovered)
    
######### RUN AN INPUT FILE IF IT WAS GIVEN AT THE COMMAND-LINE ###########
if __name__ == '__main__':
    import sys
    if len(sys.argv) != 3:
        print '\nWrong number of arguments! Do python mcsp.py <input_file> <site_budget>'
    else:
        sites, species = optimize(sys.argv[1], int(sys.argv[2]))
        print 'The near-optimal solution is to acquire sites %s, which conserves species %s, a total of %i species.' % (sites,species,species.size)